"""模型管理与推理封装
 - RoadOffsetRunner：道路分割与偏移计算（PyTorch BiSeNetV2）
 - ClassificationRunner：障碍分类（MobileViT，若依赖或权重缺失则优雅降级）
"""

import os
import sys
import time
import logging
from typing import Optional, Tuple, Dict

import numpy as np

from .config import setup_logging, MEDIA_CONFIG


class RoadOffsetRunner:
    """封装道路分割与偏移计算"""
    def __init__(self, base_dir: Optional[str] = None, weights_path: Optional[str] = None):
        self.logger = setup_logging('RoadOffset')
        self.ready = False
        try:
            base_dir = base_dir or MEDIA_CONFIG.get('road_offset_model_dir')
            if not base_dir:
                self.logger.warning('未设置 road_offset_model_dir')
                return
            # 允许相对导入 deploy 下的模块
            deploy_dir = os.path.join(base_dir)
            sys.path.append(deploy_dir)

            # 直接从deploy目录导入模块
            import inference  # type: ignore
            RoadSegmentor = inference.RoadSegmentor
            weights_path = weights_path or os.path.join(base_dir, 'models', 'weights', 'model_final.pth')
            self.runner = RoadSegmentor(weights_path=weights_path)
            self.ready = True
            self.logger.info('RoadOffsetRunner 初始化成功')
        except Exception as e:
            self.logger.error(f'RoadOffsetRunner 初始化失败: {e}')
            self.ready = False

    def infer(self, bgr_frame: np.ndarray) -> Optional[Dict]:
        if not self.ready:
            return None
        try:
            from utils.data_utils import apply_roi, calculate_road_position  # type: ignore
            road_mask, _overlay = self.runner.infer(bgr_frame, return_overlay=False)
            roi_mask = apply_roi(road_mask, top_ratio=0.4)
            ratio, offset = calculate_road_position(roi_mask)
            return {
                'road_ratio': ratio,
                'center_offset': offset
            }
        except Exception as e:
            self.logger.error(f'RoadOffset 推理失败: {e}')
            return None


class ClassificationRunner:
    """封装MobileViT障碍分类
    若缺少依赖或权重，则优雅降级返回None。
    """
    def __init__(self, base_dir: Optional[str] = None, model_path: Optional[str] = None, device: str = 'cpu'):
        self.logger = setup_logging('Classifier')
        self.ready = False
        self.device = device
        try:
            base_dir = base_dir or MEDIA_CONFIG.get('classification_model_dir')
            if not base_dir:
                self.logger.warning('未设置 classification_model_dir')
                return

            import torch  # noqa: F401
            import torchvision.transforms as transforms
            import cv2  # noqa: F401

            # MobileViT 模型定义位于部署目录的上一级（根据提供脚本）
            raspberry_dir = os.path.join(base_dir)
            parent_dir = os.path.dirname(raspberry_dir)
            sys.path.append(parent_dir)

            try:
                from mobilevit_complete import MobileViT  # type: ignore
            except Exception as e:
                self.logger.error(f'无法导入MobileViT定义: {e}')
                return

            import torch
            import torchvision.transforms as transforms

            model_path = model_path or os.path.join(base_dir, 'best_mobilevit_model.pth')
            self.class_names = ['台阶', '杆子', '正常', '路口', '路障']
            self.risk_levels = {
                '正常': 'low',
                '台阶': 'medium',
                '杆子': 'medium',
                '路口': 'high',
                '路障': 'high'
            }

            # 构造模型
            self.model = MobileViT(
                image_size=(224, 224),
                dims=[144, 192, 240],
                channels=[16, 32, 64, 64, 96, 96, 128, 128, 160, 160, 640],
                num_classes=5
            )

            # 加载权重
            if not os.path.isfile(model_path):
                self.logger.warning(f'未找到分类权重: {model_path}')
                return

            checkpoint = torch.load(model_path, map_location=torch.device(device))
            self.model.load_state_dict(checkpoint)
            self.model.to(torch.device(device))
            self.model.eval()

            # 预处理
            self.transform = transforms.Compose([
                transforms.ToPILImage(),
                transforms.Resize((224, 224)),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
            ])

            self.ready = True
            self.logger.info('ClassificationRunner 初始化成功')
        except Exception as e:
            self.logger.error(f'ClassificationRunner 初始化失败: {e}')
            self.ready = False

    def infer(self, bgr_frame: np.ndarray) -> Optional[Dict]:
        if not self.ready:
            return None
        try:
            import cv2
            import torch

            # BGR->RGB 再做标准化
            frame_rgb = cv2.cvtColor(bgr_frame, cv2.COLOR_BGR2RGB)
            tensor = self.transform(frame_rgb).unsqueeze(0)
            with torch.no_grad():
                outputs = self.model(tensor)
                probabilities = torch.softmax(outputs, dim=1)
                confidence, predicted = torch.max(probabilities, 1)

            class_idx = int(predicted.item())
            class_name = self.class_names[class_idx]
            confidence_score = float(confidence.item())
            risk_level = self.risk_levels.get(class_name, 'low')

            return {
                'class_name': class_name,
                'confidence': confidence_score,
                'risk_level': risk_level
            }
        except Exception as e:
            self.logger.error(f'分类推理失败: {e}')
            return None


class ModelManager:
    """统一管理两个模型的加载与推理"""
    def __init__(self):
        self.logger = setup_logging('Models')
        self.road = None
        self.cls = None

    def init(self):
        try:
            self.road = RoadOffsetRunner(
                base_dir=MEDIA_CONFIG.get('road_offset_model_dir')
            )
        except Exception as e:
            self.logger.error(f'初始化道路偏移模型失败: {e}')
            self.road = None

        try:
            self.cls = ClassificationRunner(
                base_dir=MEDIA_CONFIG.get('classification_model_dir'),
                model_path=os.path.join(
                    MEDIA_CONFIG.get('classification_model_dir', ''),
                    'best_mobilevit_model.pth'
                )
            )
        except Exception as e:
            self.logger.error(f'初始化分类模型失败: {e}')
            self.cls = None

    def infer_both(self, bgr_frame: np.ndarray) -> Dict:
        result: Dict = {}
        if self.road:
            road_res = self.road.infer(bgr_frame)
            if road_res:
                result['road'] = road_res
        if self.cls:
            cls_res = self.cls.infer(bgr_frame)
            if cls_res:
                result['classification'] = cls_res
        return result